﻿# utils/plotly_plot.py
import plotly.graph_objects as go
import pandas as pd
import json

def plot_kline_to_json(df, buy_signals=None, sell_signals=None):
    fig = go.Figure()

    # 添加K线
    fig.add_trace(go.Candlestick(
        x=df.index,
        open=df['open'],
        high=df['high'],
        low=df['low'],
        close=df['close'],
        name="K线"
    ))

    # 处理买入信号
    if buy_signals is not None and not buy_signals.empty:
        try:
            signal_times = pd.to_datetime(buy_signals['time'])
            matched_lows = df.loc[df.index.isin(signal_times), 'low']
            if not matched_lows.empty:
                fig.add_trace(go.Scatter(
                    x=matched_lows.index,
                    y=matched_lows * 0.99,
                    mode="markers",
                    marker=dict(symbol="triangle-up", size=10, color="green"),
                    name="买入"
                ))
        except Exception as e:
            print(f"⚠️ 绘制买入信号失败: {e}")

    # 处理卖出信号
    if sell_signals is not None and not sell_signals.empty:
        try:
            signal_times = pd.to_datetime(sell_signals['time'])
            matched_highs = df.loc[df.index.isin(signal_times), 'high']
            if not matched_highs.empty:
                fig.add_trace(go.Scatter(
                    x=matched_highs.index,
                    y=matched_highs * 1.01,
                    mode="markers",
                    marker=dict(symbol="triangle-down", size=10, color="red"),
                    name="卖出"
                ))
        except Exception as e:
            print(f"⚠️ 绘制卖出信号失败: {e}")

    # 更新布局
    fig.update_layout(
        title="量化回测动画",
        xaxis_title="时间",
        yaxis_title="价格",
        xaxis_rangeslider_visible=False,
        template="plotly_dark",  # 只用字符串
        height=600,
        margin=dict(l=20, r=20, t=40, b=20),
        autosize=True
    )

    # ✅ 安全序列化 traces
    traces = []
    for trace in fig.data:
        trace_dict = {}

        if hasattr(trace, 'x'):
            trace_dict['x'] = [str(x) for x in trace.x]

        if hasattr(trace, 'name'):
            trace_dict['name'] = trace.name
        if hasattr(trace, 'type'):
            trace_dict['type'] = trace.type
        if hasattr(trace, 'mode'):
            trace_dict['mode'] = trace.mode
        if hasattr(trace, 'marker'):
            # ✅ 确保 marker 可序列化
            trace_dict['marker'] = trace.marker.to_plotly_json()

        # 处理不同 trace 类型
        if isinstance(trace, go.Candlestick):
            trace_dict['open'] = trace.open.tolist()
            trace_dict['high'] = trace.high.tolist()
            trace_dict['low'] = trace.low.tolist()
            trace_dict['close'] = trace.close.tolist()
        elif hasattr(trace, 'y'):
            trace_dict['y'] = trace.y.tolist() if hasattr(trace.y, 'tolist') else trace.y

        traces.append(trace_dict)

    # ✅ 手动构建 layout，避免包含不可序列化的对象
    layout = {
        'title': fig.layout.title.text,
        'xaxis': {
            'title': fig.layout.xaxis.title.text,
            'rangeslider': {'visible': False}
        },
        'yaxis': {'title': fig.layout.yaxis.title.text},
        'template': 'plotly_dark',  # ✅ 只保留字符串
        'height': fig.layout.height,
        'margin': {
            'l': fig.layout.margin.l,
            'r': fig.layout.margin.r,
            't': fig.layout.margin.t,
            'b': fig.layout.margin.b
        }
    }

    # ✅ 最终序列化
    return json.dumps({'traces': traces, 'layout': layout}, ensure_ascii=False, indent=2)